package org.example.offical.doc.ai.service;

import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.content.aggregator.DefaultContentAggregator;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.service.V;
import org.example.offical.doc.ModelUtils;

/**
 * @author superMan
 * @since fish_temp_since
 */
public class AiServiceRAG {
    public static void main(String[] args) {

        // 基于redis 的向量库
        RedisEmbeddingStore redisEmbeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .indexName("ai-service-rag")
                .prefix("rag:")
                .dimension(1536)
                .build();
        // openAI 向量模型
        EmbeddingModel openAiEmbeddingDemoModel = ModelUtils.getOpenAiEmbeddingDemoModel();

        redisEmbeddingStore.add(openAiEmbeddingDemoModel.embed("李四的bmi是25").content(), TextSegment.textSegment("李四的bmi是25"));
        redisEmbeddingStore.add(openAiEmbeddingDemoModel.embed("李四的身高是181").content(), TextSegment.textSegment("李四的身高是181"));
        redisEmbeddingStore.add(openAiEmbeddingDemoModel.embed("李四的体重是90KG").content(), TextSegment.textSegment("李四的体重是90KG"));
        redisEmbeddingStore.add(openAiEmbeddingDemoModel.embed("张三的年龄是25").content(), TextSegment.textSegment("张三的年龄是25"));
        redisEmbeddingStore.add(openAiEmbeddingDemoModel.embed("王二的bmi是25").content(), TextSegment.textSegment("王二的bmi是25"));

        EmbeddingStoreContentRetriever embeddingStoreContentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(redisEmbeddingStore)
                .embeddingModel(openAiEmbeddingDemoModel)
                .maxResults(5)
                .minScore(0.6)
                .build();

        // System.out.println("召回测试——-----");
        embeddingStoreContentRetriever.retrieve(Query.from("李四的bmi")).forEach(System.out::println);
        // System.out.println("召回测试——-----");

        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(ModelUtils.getOpenAiDemoModel())
                .retrievalAugmentor(DefaultRetrievalAugmentor.builder()
                        // 文本召回
                        .contentRetriever(embeddingStoreContentRetriever)
                        // 重排序
                        .contentAggregator(new DefaultContentAggregator())
                        // 注入
                        .contentInjector(DefaultContentInjector.builder()
                                // prompt
                                .promptTemplate(PromptTemplate.from("""
                                        你是健身专家，请回答用户的问题：{{it}}
                                        """))
                                .build())
                        .build())
                .build();


        System.out.println(assistant.chat("请问李四的BMI是多少"));
    }


    interface Assistant {
        String chat(@V("it") String userMsg);
    }
}
